#!/usr/bin/env python3
"""
Systemtap MCP Server

面向 systemtap 软件及其相关服务（stap-exporter、stap-server、systemtap）的命令行操作。
提供对 systemtap 相关服务的启动、停止、重启、状态检查等系统服务管理功能。

MCP tools 列表：
1. stap_exporter_service_control - 控制 stap-exporter 服务
   - 参数：action (str): 操作类型，支持 'start', 'stop', 'restart', 'status'
   - 返回：统一 JSON 结构

2. stap_server_service_control - 控制 stap-server 服务  
   - 参数：action (str): 操作类型，支持 'start', 'stop', 'restart', 'status'
   - 返回：统一 JSON 结构

3. systemtap_service_control - 控制 systemtap 服务
   - 参数：action (str): 操作类型，支持 'start', 'stop', 'restart', 'status', 'reload'
   - 返回：统一 JSON 结构
"""

import subprocess
import json
from mcp.server.fastmcp import FastMCP, Context

mcp = FastMCP("Systemtap MCP Server")

@mcp.tool()
def stap_exporter_service_control(action: str) -> dict:
    """
    控制 stap-exporter 服务的启动、停止、重启和状态检查
    
    参数:
        action (str): 操作类型，支持 'start', 'stop', 'restart', 'status'
    
    返回:
        dict: 统一 JSON 结构
        {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误
        }
    """
    valid_actions = ['start', 'stop', 'restart', 'status']
    if action not in valid_actions:
        return {
            "success": False,
            "command": f"systemctl {action} stap-exporter.service",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Invalid action: {action}. Supported actions: {', '.join(valid_actions)}"
        }
    
    cmd = ["systemctl", action, "stap-exporter.service"]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 30 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

@mcp.tool()
def stap_server_service_control(action: str) -> dict:
    """
    控制 stap-server 服务的启动、停止、重启和状态检查
    
    参数:
        action (str): 操作类型，支持 'start', 'stop', 'restart', 'status'
    
    返回:
        dict: 统一 JSON 结构
        {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误
        }
    """
    valid_actions = ['start', 'stop', 'restart', 'status']
    if action not in valid_actions:
        return {
            "success": False,
            "command": f"systemctl {action} stap-server.service",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Invalid action: {action}. Supported actions: {', '.join(valid_actions)}"
        }
    
    cmd = ["systemctl", action, "stap-server.service"]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 30 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

@mcp.tool()
def systemtap_service_control(action: str) -> dict:
    """
    控制 systemtap 服务的启动、停止、重启、状态检查和重载
    
    参数:
        action (str): 操作类型，支持 'start', 'stop', 'restart', 'status', 'reload'
    
    返回:
        dict: 统一 JSON 结构
        {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误
        }
    """
    valid_actions = ['start', 'stop', 'restart', 'status', 'reload']
    if action not in valid_actions:
        return {
            "success": False,
            "command": f"systemctl {action} systemtap.service",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Invalid action: {action}. Supported actions: {', '.join(valid_actions)}"
        }
    
    cmd = ["systemctl", action, "systemtap.service"]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 30 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

if __name__ == "__main__":
    mcp.run()